import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter
import matplotlib
# import matplotlib.pyplot as plt
# plt.rcParams['font.sans-serif']=['KaiTi']
# plt.rcParams['axes.unicode_minus'] = False


class FlightDataLogger:
    def __init__(self):
        # 初始化数据记录结构
        self.flight_data = {
            'episode': [],
            'step': [],
            'delta_long': [],  # 经度偏差 (m)
            'delta_lat': [],  # 纬度偏差 (m)
            'altitude': [],  # 当前高度 (m)
            'altitude_deviation': [],  # 高度偏差 (m)
            'path_deviation': [],  # 路径偏差 (m)
            'heading_error': [],  # 航向误差 (度)
            'IAS': [],  # 指示空速 (m/s)
            # 新增奖励参数
            'dist_reward': [],
            'heading_reward': [], # 航向奖励
            'speed_reward': [], # 速度奖励
            'vspeed_reward': [], # 垂直速度奖励
            'stability_reward': [],
            'total_reward': []
        }

    def log_data(self, episode, step, state, reward_data=None):
        """记录飞行数据"""
        self.flight_data['episode'].append(episode)
        self.flight_data['step'].append(step)
        self.flight_data['delta_long'].append(state[20])  # 经度偏差
        self.flight_data['delta_lat'].append(state[21])  # 纬度偏差
        self.flight_data['altitude'].append(state[32])  # 当前高度 (LLA_balt)
        self.flight_data['altitude_deviation'].append(state[22])  # 高度偏差
        self.flight_data['path_deviation'].append(state[28])  # 路径偏差
        # 计算航向误差 (当前航向与目标航向之间的最小角度)
        heading_error = min(abs(state[2] - state[27]), 360 - abs(state[2] - state[27]))
        self.flight_data['heading_error'].append(heading_error)  # 航向偏差
        self.flight_data['IAS'].append(state[1])  # 指示空速

        # 记录奖励数据
        if reward_data:
            for key in ['dist_reward', 'heading_reward', 'speed_reward',
                        'vspeed_reward', 'stability_reward', 'total_reward']:
                self.flight_data[key].append(reward_data.get(key, 0))
        else:
            for key in ['dist_reward', 'heading_reward', 'speed_reward',
                        'vspeed_reward', 'stability_reward', 'total_reward']:
                self.flight_data[key].append(0)

    def save_to_csv(self, filename="results//flight_parameters.csv"):
        """将数据保存到 CSV 文件"""
        df = pd.DataFrame(self.flight_data)
        df.to_csv(filename, index=False)
        print(f"飞行数据已保存到 {filename}")

    def plot_parameters(self):
        """将所有参数绘制成二维图表"""
        df = pd.DataFrame(self.flight_data)

        # 设置科学记数法显示
        formatter = ScalarFormatter(useMathText=True)
        formatter.set_scientific(True)
        formatter.set_powerlimits((-3, 4))

        plt.figure(figsize=(20, 20))

        # 1. 经度和纬度偏差
        plt.subplot(4, 3, 1)
        plt.plot(df['episode'], df['delta_long'], 'r-', label='Longitude dev (m)')
        plt.plot(df['episode'], df['delta_lat'], 'b-', label='Latitude dev (m)')
        plt.xlabel('episode')
        plt.ylabel('Deviation (m)')
        plt.title('Longitude/Latitude Deviation vs Time')
        plt.legend()
        plt.gca().yaxis.set_major_formatter(formatter)
        plt.grid(True)

        # 2. 高度和高度偏差
        plt.subplot(4, 3, 2)
        plt.plot(df['episode'], df['altitude'], 'g-', label='Altitude (m)')
        plt.plot(df['episode'], df['altitude_deviation'], 'm-', label='Alt deviation (m)')
        plt.xlabel('episode')
        plt.ylabel('Altitude (m)')
        plt.title('Altitude and Deviation vs Time')
        plt.legend()
        plt.gca().yaxis.set_major_formatter(formatter)
        plt.grid(True)

        #  距离奖励
        plt.subplot(4, 3, 3)
        plt.plot(df['episode'], df['dist_reward'], 'b-', label='Distance Reward')
        plt.xlabel('episode')
        plt.ylabel('Reward')
        plt.title('Distance Reward vs Time')
        plt.grid(True)

        # 3. 路径偏差
        plt.subplot(4, 3, 4)
        plt.plot(df['episode'], df['path_deviation'], 'c-')
        plt.xlabel('episode')
        plt.ylabel('Path deviation (m)')
        plt.title('Path Deviation vs Time')
        plt.gca().yaxis.set_major_formatter(formatter)
        plt.grid(True)

        # 4. 航向误差
        plt.subplot(4, 3, 5)
        plt.plot(df['episode'], df['heading_error'], 'y-')
        plt.xlabel('episode')
        plt.ylabel('Heading error (deg)')
        plt.title('Heading Error vs Time')
        plt.grid(True)

        #  航向奖励
        plt.subplot(4, 3, 6)
        plt.plot(df['episode'], df['heading_reward'], 'g-', label='Heading Reward')
        #plt.plot(df['episode'], df['speed_reward'], 'r-', label='Speed Reward')
        #plt.plot(df['episode'], df['vspeed_reward'], 'm-', label='VSpeed Reward')
        plt.xlabel('episode')
        plt.ylabel('Reward')
        plt.title('Heading Rewards vs Time')
        plt.legend()
        plt.grid(True)

        # 5. 空速
        plt.subplot(4, 3, 7)
        plt.plot(df['episode'], df['IAS'], 'k-')
        plt.xlabel('episode')
        plt.ylabel('IAS (m/s)')
        plt.title('Indicated Airspeed vs Time')
        plt.grid(True)

        # 速度奖励
        plt.subplot(4, 3, 8)
        # plt.plot(df['episode'], df['heading_reward'], 'g-', label='Heading Reward')
        plt.plot(df['episode'], df['speed_reward'], 'r-', label='Speed Reward')
        plt.plot(df['episode'], df['vspeed_reward'], 'm-', label='VSpeed Reward')
        plt.xlabel('episode')
        plt.ylabel('Reward')
        plt.title('Speed/VSpeed Rewards vs Time')
        plt.legend()
        plt.grid(True)

        # 8. 总奖励和稳定性奖励
        plt.subplot(4, 3, 9)
        plt.plot(df['episode'], df['total_reward'], 'k-', label='Total Reward')
        plt.plot(df['episode'], df['stability_reward'], 'c-', label='Stability Reward')
        plt.xlabel('episode')
        plt.ylabel('Reward')
        plt.title('Total and Stability Rewards vs Time')
        plt.legend()
        plt.grid(True)




        plt.tight_layout()
        # plt.show()
        plt.savefig('results//flight_parameters.png')
        plt.close()
        print("飞行参数图表已保存为 flight_parameters.png")